# Code for gaussian radial basis function.

import jax
import jax.numpy as jnp
import numpy as np
import haiku as hk

from typing import Optional, Union, Tuple, List
from jax import nn as nn
from .rbf_base import RadialBasisFunctions, _rbf_register

@_rbf_register('gaussian')
class GaussianBasis(RadialBasisFunctions):
    r"""Gaussian type RBF.
    Args:
        r_max (Length):         Maximum distance. Defatul: 1 nm.
        r_min (Length):         Minimum distance. Default: 0 nm.
        sigma (float):          Simga. Default: 0.03 nm.
        delta (float):          Space interval. Default: 0.016 nm.
        num_basis (int):        Number of basis functions. Defatul: None.
        clip_distance (bool):   Whether to clip the value of distance. Default: False.
        name (str):             Name of the module. Default: "gaussian_basis".
    """

    def __init__(self,
                 fp_type = jnp.float32,
                 r_max: float = 1.0,
                 r_min: float = 0.0,
                 sigma: float = 0.03,
                 delta: float = 0.016,
                 num_basis: Optional[int] = None,
                 clip_distance: bool = False,
                 name: str = "gaussian_basis",
                 ):
        super().__init__(fp_type=fp_type,
                         r_max=r_max,
                         r_min=r_min,
                         num_basis=num_basis,
                         clip_distance=clip_distance,
                         name=name)

        if r_max <= r_min:
            raise ValueError("[utils/rbf/GaussianBasis] r_max should be larger than r_min.")
        
        self.r_range = self.r_max - self.r_min
        
        if num_basis is None and delta is None:
            raise TypeError('[utils/rbf/GaussianBasis] "num_basis" and "delta" cannot both be "None".')
        if num_basis is not None and num_basis <= 0:
            raise ValueError('[utils/rbf/GaussianBasis] "num_basis" must be larger than 0.')
        if delta is not None and delta <= 0:
            raise ValueError('[utils/rbf/GaussianBasis] "delta" must be larger than 0.')
        
        self.delta = delta
        self.sigma = sigma
        self.coefficient = -0.5 * jnp.reciprocal(jnp.square(self.sigma))

        if self.delta is None and self.num_basis is not None:
            self.offset = jnp.linspace(self.r_min, self.r_max, self.num_basis, dtype=self.fp_type)
            self.delta = self.r_range / (self.num_basis - 1)
        else:
            ## Liyh: trouble here
            if self.num_basis is None:
                self.num_basis = np.ceil(self.r_range / self.delta) + 1
                self.num_basis = int(self.num_basis)
            self.offsets =  self.r_min + jnp.arange(0, self.num_basis, dtype=self.fp_type) * self.delta ## numpy
        
    def __call__(self, distance: jnp.ndarray) -> jnp.ndarray:
        r"""Compute gaussian type RBF.

        Args: 
            distance (Array):                 Distance matrix. Shape: (A, A).
        Returns: 
            radial basis embedding (Array):   Embedding of distance matrix. Shape: (A, A, num_basis).

        """

        if self.clip_distance:
            distance = jnp.clip(distance, self.r_min, self.r_max)
        
        # (..., ) -> (..., 1)
        distance = jnp.expand_dims(distance, axis=-1)
        # (..., 1) - (..., num_basis) -> (..., num_basis)
        diff = distance - self.offsets
        # (..., num_basis) -> (..., num_basis)
        rbf = jnp.exp(self.coefficient * jnp.square(diff))

        return rbf

    def __str__(self) -> str:
        return 'GaussianBasis<>'
